import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
import lightgbm as lgb
import matplotlib.pyplot as plt
from datetime import timedelta
# from sklearn.impute import SimpleImputer  # 替换为SimpleImputer

# ======================
# 1. 模拟数据生成（实际使用时替换为真实数据）
# ======================
# def generate_dummy_data(start_date='2023-01-01', end_date='2023-12-31', sku_list=['SKU_A', 'SKU_B']):
#     date_range = pd.date_range(start_date, end_date)
#     data = []
#
#     for sku in sku_list:
#         base_sales = np.random.randint(50, 150)
#         for date in date_range:
#             # 生成带有季节性和随机性的销量
#             seasonality = 20 * np.sin(2 * np.pi * (date.dayofyear / 365))
#             noise = np.random.normal(0, 10)
#             sales = max(0, int(base_sales + seasonality + noise))
#
#             data.append({
#                 'date_id': date,
#                 'sku': sku,
#                 'sales': sales
#             })
#
#     return pd.DataFrame(data)


# 生成模拟数据（替换为你的真实数据加载）
df = pd.read_csv('2024-LO2312019BFFFS.csv', low_memory=False, encoding='GBK')
df['date_id'] = pd.to_datetime(df['date_id'],format='%Y%m%d')
df = df.sort_values(['sku', 'date_id']).reset_index(drop=True)

plt.rcParams['font.sans-serif'] = ['SimHei']
# ======================
# 2. 特征工程
# ======================
def create_features(df):
    # 时间特征
    df['year'] = df['date_id'].dt.year
    df['month'] = df['date_id'].dt.month
    df['day'] = df['date_id'].dt.day
    df['dayofweek'] = df['date_id'].dt.dayofweek
    df['weekofyear'] = df['date_id'].dt.isocalendar().week.astype(int)

    # 滞后特征
    df['lag_3'] = df.groupby('sku')['sales'].shift(3)
    df['lag_7'] = df.groupby('sku')['sales'].shift(7)

    # 滚动特征
    df['rolling_7_mean'] = df.groupby('sku')['sales'].transform(
        lambda x: x.rolling(7, min_periods=1).mean()
    )
    df['rolling_14_std'] = df.groupby('sku')['sales'].transform(
        lambda x: x.rolling(30, min_periods=1).std()
    )

    # 处理缺失值
    df.ffill()
    df['sku'] = df['sku'].astype('category')
    df['sku'] = df['sku'].astype('category')

    return df.drop('date_id', axis=1)


# 执行特征工程
featured_df = create_features(df.copy())

# ======================
# 3. 数据分割
# ======================
train_size = int(len(featured_df) * 0.8)
train = featured_df.iloc[:train_size]
test = featured_df.iloc[train_size:]

X_train = train.drop('sales', axis=1)
y_train = train['sales']
X_test = test.drop('sales', axis=1)
y_test = test['sales']

# ======================
# 4. 模型训练
# ======================
params = {
    'objective': 'regression',
    'metric': 'mae',
    'boosting_type': 'gbdt',
    'learning_rate': 0.05,
    'num_leaves': 31,
    'min_data_in_leaf': 20,
    'feature_fraction': 0.8,
    'verbosity': -1,
    'seed': 42
}

# my_imputer = SimpleImputer(strategy='mean')  # 使用SimpleImputer
# X_train = my_imputer.fit_transform(X_train)
# X_test = my_imputer.transform(X_test)

train_data = lgb.Dataset(X_train, label=y_train, categorical_feature=['sku'])
model = lgb.train(
    params,
    train_data,
    num_boost_round=1000,
    valid_sets=[train_data],
    callbacks=[lgb.early_stopping(stopping_rounds=50)]
)

# ======================
# 5. 模型评估
# ======================
y_pred = model.predict(X_test)
print(f"\n模型评估结果:")
print(f"MAE: {mean_absolute_error(y_test, y_pred):.2f}")
print(f"RMSE: {np.sqrt(mean_squared_error(y_test, y_pred)):.2f}")


# ======================
# 6. 未来预测函数
# ======================
def predict_future_sales(model, original_df, sku_id, forecast_days=7):
    """预测指定SKU未来N天的销量"""
    df = original_df.copy()
    sku_data = df[df['sku'] == sku_id].sort_values('date_id')

    # 生成未来日期
    last_date = sku_data['date_id'].max()
    future_dates = [last_date + timedelta(days=i) for i in range(1, forecast_days + 1)]

    # 构建未来数据框架
    future_df = pd.DataFrame({
        'date_id': future_dates,
        'sku': sku_id,
        'sales': np.nan
    })

    full_df = pd.concat([sku_data, future_df], ignore_index=True)

    # 递归预测
    for i in range(len(sku_data), len(full_df)):
        # 生成临时特征
        temp_df = create_features(full_df.iloc[:i + 1].copy())

        # 获取最新特征
        latest_features = temp_df.iloc[-1:].drop('sales', axis=1)

        # 预测并更新
        pred = max(0, model.predict(latest_features)[0])
        full_df.loc[i, 'sales'] = pred

        # 更新原始数据用于下次迭代
        full_df.loc[i, 'sales'] = pred

    return full_df[['date_id', 'sku', 'sales']].tail(forecast_days)


# ======================
# 7. 执行预测示例
# ======================
# 选择第一个SKU进行演示
# target_sku = df['sku'].unique()[0]

# target_sku = 'LO2312019AXNHS'
target_sku = 'LO2312019BFFFS'

# 获取历史数据
history = df[df['sku'] == target_sku][['date_id', 'sales']]

# 执行预测
predictions = predict_future_sales(model, df, target_sku)

print("\n历史销量最后7天:")
print(history.tail(7))
print("\n未来7天预测:")
print(predictions)

# ======================
# 8. 结果可视化
# ======================
plt.figure(figsize=(12, 6))
plt.plot(history['date_id'], history['sales'], label='历史销量', marker='o')
plt.plot(predictions['date_id'], predictions['sales'], label='预测销量',
         linestyle='--', marker='x', color='red')
plt.title(f'SKU {target_sku} 销量预测')
plt.xlabel('日期')
plt.ylabel('销量')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()